import os
import torchvision.datasets as dset


def main():
    # ==========================================
    # 👇 在这里修改为你想要的绝对路径
    # 例如 Windows: r"D:\Datasets"
    # 例如 Mac/Linux: "/Users/lijiahao/Documents/MyData"
    custom_download_path = "/Users/lijiahao/Documents/MyData"
    # ==========================================

    # 1. 确保目录存在
    if not os.path.exists(custom_download_path):
        os.makedirs(custom_download_path)
        print(f"创建下载目录: {custom_download_path}")

    print(f"准备将数据下载到: {custom_download_path}")

    # 2. 开始下载 (root 参数指定了下载位置)
    print("\n[1/3] 正在检查/下载 VOC 2007 TrainVal...")
    dset.VOCDetection(root=custom_download_path, year='2007', image_set='trainval', download=True)

    print("\n[2/3] 正在检查/下载 VOC 2007 Test...")
    dset.VOCDetection(root=custom_download_path, year='2007', image_set='test', download=True)

    print("\n[3/3] 正在检查/下载 VOC 2012 TrainVal...")
    dset.VOCDetection(root=custom_download_path, year='2012', image_set='trainval', download=True)

    # ---------------------------------------------------------
    # 关键步骤：建立软链接
    # 让 faster-rcnn 代码以为数据就在项目里，实际上数据在你的 custom_download_path
    # ---------------------------------------------------------

    # 获取当前脚本所在目录（假设脚本放在项目根目录）
    project_root = os.path.dirname(os.path.abspath(__file__))
    project_data_dir = os.path.join(project_root, 'data')  # 项目里的 ./data

    # torchvision 会在你的自定义路径下生成 'VOCdevkit' 文件夹
    real_data_location = os.path.join(custom_download_path, 'VOCdevkit')

    # 我们要在项目里生成的链接位置: ./data/VOCdevkit
    link_destination = os.path.join(project_data_dir, 'VOCdevkit')
    link_destination_2007 = os.path.join(project_data_dir, 'VOCdevkit2007')

    # 确保项目里的 data 目录存在
    if not os.path.exists(project_data_dir):
        os.makedirs(project_data_dir)

    print("\n正在配置软链接，以便训练代码能找到数据...")

    # 建立链接 1: data/VOCdevkit -> /你的绝对路径/VOCdevkit
    if not os.path.exists(link_destination):
        try:
            os.symlink(real_data_location, link_destination)
            print(f"链接成功: {link_destination} -> {real_data_location}")
        except OSError as e:
            print(f"创建链接失败: {e}")
    else:
        print(f"链接已存在: {link_destination}")

    # 建立链接 2: data/VOCdevkit2007 -> /你的绝对路径/VOCdevkit (为了兼容性)
    if not os.path.exists(link_destination_2007):
        try:
            os.symlink(real_data_location, link_destination_2007)
            print(f"链接成功: {link_destination_2007} -> {real_data_location}")
        except OSError as e:
            print(f"创建链接失败: {e}")

    print("\n所有操作完成！数据实际存储在:", custom_download_path)


if __name__ == '__main__':
    main()